#!/usr/bin/python3
# -*- coding: utf-8 -*-
import os
import struct

base_path = r'/rootfs/media/kasim/DataSet/FaceVideo'
out_path = r'/rootfs/media/kasim/DataSet/FaceVideo'
video_list_bbox_count_files = [
    r'/rootfs/media/kasim/DataSet/RecordEx/Face_video_list_bbox_count.txt',
    r'/rootfs/media/kasim/DataSet/Record/Face_video_list_bbox_count.txt',
]
video_list_bbox_files = [
    r'/rootfs/media/kasim/DataSet/RecordEx/Face_video_list_bbox.dat',
    r'/rootfs/media/kasim/DataSet/Record/Face_video_list_bbox.dat',
]
video_list_filter_file = r'/rootfs/media/kasim/DataSet/FaceVideo/file_list.txt'

select_class = 'Face'

select_class_label = {
    'Person': 0,
    'Cat': 1,
    'Dog': 2,
    'BabyCar': 3,
    'Face': 4,
}
select_label = select_class_label[select_class]

THRESHOLDS = [
    0.3,  # Person
    0.3,  # Cat
    0.3,  # Dog
    0.3,  # BabyCar
    0.3,  # Face
]


def main():
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    os.system('chmod a+wr {}'.format(out_path))

    out_list_file_path = os.path.join(out_path, '{}_video_list_bbox_count.txt'.format(select_class.lower()))
    out_bbox_path = os.path.join(out_path, '{}_video_list_bbox.dat'.format(select_class.lower()))

    out_list_file = open(out_list_file_path, 'w')
    out_bbox_file = open(out_bbox_path, 'wb')

    try:
        filter_video_dict = set()
        with open(video_list_filter_file, 'r') as file:
            for video_info in file.readlines():
                video_name = video_info.split()[0]
                filter_video_dict.add(video_name)
        print('Filter File Count: {}'.format(len(filter_video_dict)))

        score_thr = THRESHOLDS[select_label]
        video_bboxes_map = {}
        file_count = 0
        for video_list_bbox_count_file, video_list_bbox_file in zip(video_list_bbox_count_files, video_list_bbox_files):
            with open(video_list_bbox_count_file, 'r') as file:
                with open(video_list_bbox_file, 'rb') as bbox_file:
                    for video_info in file.readlines():
                        video_info = video_info.split(',')
                        video_name = video_info[0].strip()
                        BBox_count = int(video_info[1].strip())

                        file_count += 1
                        if file_count % 1000 == 0:
                            print('Read File BBox Count: {}'.format(file_count))

                        if video_name in video_bboxes_map.keys():
                            bbox_info = bbox_file.read(28*BBox_count)
                            if len(bbox_info) < 28*BBox_count:
                                break
                            continue
                        else:
                            bbox_infos = []
                            for j in range(BBox_count):
                                bbox_info = bbox_file.read(28)
                                if len(bbox_info) < 28:
                                    break
                                bbox = struct.unpack('6i1f', bbox_info)
                                label = bbox[1]
                                if select_label != label:
                                    continue
                                # frame_id = bbox[0]
                                score = bbox[6]
                                if score < score_thr:
                                    continue
                                bbox_infos.append(bbox_info)
                            if len(bbox_infos) > 0:
                                video_bboxes_map[video_name] = bbox_infos
        print('Read File BBox Count Finish: {}'.format(file_count))

        file_count = 0
        total_file_count = len(video_bboxes_map)
        for video_name, bbox_infos in video_bboxes_map.items():
            file_count += 1
            if file_count % 1000 == 0:
                print('File Count: {}/{}'.format(file_count, total_file_count))
            if video_name not in filter_video_dict:
                continue

            bbox_count = len(bbox_infos)
            if bbox_count > 0:
                file_info_str = '{},{}\n'.format(video_name, bbox_count)
                out_list_file.write(file_info_str)
                for bbox_info in bbox_infos:
                    out_bbox_file.write(bbox_info)

    except Exception as e:
        print(e)
    print('File Count: {}/{}'.format(file_count, total_file_count))

    out_list_file.close()
    out_bbox_file.close()
    os.system('chmod a+wr {}'.format(out_list_file_path))
    os.system('chmod a+wr {}'.format(out_bbox_path))

    print('finish!!!!')


if __name__ == '__main__':
    main()
